import rclpy
from rclpy.node import Node
from nav_msgs.msg import Odometry
from unitree_go.msg import SportModeState, LowState
from tf2_ros import TransformBroadcaster 
from geometry_msgs.msg import TransformStamped
from sensor_msgs.msg import JointState

class Driver(Node):
    def __init__(self):
        super().__init__("driver")
        self.declare_parameter("odom_frame", "odom")
        self.declare_parameter("base_frame", "base")
        self.declare_parameter("publish_tf", "true")
        self.declare_parameter("tf_publish_rate", 10.0)  # 默认10Hz

        self.odom_frame = self.get_parameter("odom_frame").value
        self.base_frame = self.get_parameter("base_frame").value
        self.publish_tf = self.get_parameter("publish_tf").value
        self.tf_publish_rate = self.get_parameter("tf_publish_rate").value

        # 存储最新的里程计数据
        self.latest_odom = None
        self.latest_transform = None

        # 发布 odom
        self.odom_pub = self.create_publisher(Odometry, "odom", 10)
        
        # 订阅运动模式状态
        self.mode_sub = self.create_subscription(
            SportModeState,
            "/lf/sportmodestate",
            self.sport_mode_state_callback,
            10
        )

        # TF 广播器
        self.tf_broadcaster = TransformBroadcaster(self)

        # 发布关节状态
        self.joint_pub = self.create_publisher(JointState, "joint_states", 10)
        
        # 订阅低层状态
        self.state_sub = self.create_subscription(
            LowState,
            "/lf/lowstate",
            self.low_state_callback,
            10
        )

        # 创建定时器，定时发布 TF
        self.tf_timer = self.create_timer(
            1.0 / self.tf_publish_rate,  # 定时周期
            self.publish_tf_callback
        )

    def low_state_callback(self, state: LowState):
        joint_state = JointState()
        joint_state.header.stamp = self.get_clock().now().to_msg()
        joint_state.name = [
            "FL_hip_joint", "FL_thigh_joint", "FL_calf_joint",
            "FR_hip_joint", "FR_thigh_joint", "FR_calf_joint",
            "RL_hip_joint", "RL_thigh_joint", "RL_calf_joint",
            "RR_hip_joint", "RR_thigh_joint", "RR_calf_joint"
        ]
        
        for i in range(12):
            q = float(state.motor_state[i].q)
            joint_state.position.append(q)

        self.joint_pub.publish(joint_state)

    def sport_mode_state_callback(self, mode: SportModeState):
        # 解析 go2 的 odom 并发布 AGIROS 的 odom
        odom = Odometry()
        odom.header.stamp = self.get_clock().now().to_msg()
        odom.header.frame_id = self.odom_frame
        odom.child_frame_id = self.base_frame

        # 位置
        odom.pose.pose.position.x = float(mode.position[0])
        odom.pose.pose.position.y = float(mode.position[1])
        odom.pose.pose.position.z = float(mode.position[2])
 
        # 线速度
        odom.twist.twist.linear.x = float(mode.velocity[0])
        odom.twist.twist.linear.y = float(mode.velocity[1])
        odom.twist.twist.linear.z = float(mode.velocity[2])

        # 角速度
        odom.twist.twist.angular.z = float(mode.yaw_speed)

        # 姿态
        odom.pose.pose.orientation.w = float(mode.imu_state.quaternion[0])
        odom.pose.pose.orientation.x = float(mode.imu_state.quaternion[1])
        odom.pose.pose.orientation.y = float(mode.imu_state.quaternion[2])
        odom.pose.pose.orientation.z = float(mode.imu_state.quaternion[3])

        self.odom_pub.publish(odom)
        self.latest_odom = odom  # 存储最新的 odom 数据

        # 更新 TF 变换
        if self.publish_tf:
            transform = TransformStamped()
            transform.header.stamp = odom.header.stamp
            transform.header.frame_id = self.odom_frame
            transform.child_frame_id = self.base_frame

            transform.transform.translation.x = odom.pose.pose.position.x
            transform.transform.translation.y = odom.pose.pose.position.y
            transform.transform.translation.z = odom.pose.pose.position.z

            transform.transform.rotation = odom.pose.pose.orientation
            self.latest_transform = transform  # 存储最新的变换

    def publish_tf_callback(self):
        """定时发布 TF 变换"""
        if self.publish_tf and self.latest_transform:
            self.tf_broadcaster.sendTransform(self.latest_transform)

def main():
    rclpy.init()
    rclpy.spin(Driver())
    rclpy.shutdown()

if __name__ == "__main__":
    main()